/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef PARBSR_ILUAMG_PRECONDITIONER_H
#define PARBSR_ILUAMG_PRECONDITIONER_H

#include "parbsr_preconditioner.h"
#include "parcsr_matrix.h"

namespace YHAMG
{

class ParBSR_ILUAMG_Preconditioner : public ParBSR_Preconditioner
{
private:
	int num_levels;
	ParBSR_Matrix* A;
	ParCSR_Matrix* P;
	ParCSR_Matrix* R;
	class Smoother;
	Smoother* S;

	Par_Vector* r;
	Par_Vector* e;
	Par_Vector* g;

	void Strength(const ParCSR_Matrix& A, bool* LocLoc_s, bool* LocExt_s) const;
	void Coarsening(const ParCSR_Matrix& A, const bool* LocLoc_s, const bool* LocExt_s, int* CF_map, int AGGRESSIVE = 0) const;
	void Interpolation(const ParCSR_Matrix& A, const bool* LocLoc_s, const bool* LocExt_s, int* CF_map, ParCSR_Matrix& P, int AGGRESSIVE = 0) const;
	void RAP(const ParCSR_Matrix& R, const ParBSR_Matrix& A, const ParCSR_Matrix& P, ParBSR_Matrix& C) const;
	void Sparsification(const ParBSR_Matrix& A, ParBSR_Matrix& B) const;
	void Smooth(int level, const Par_Vector& b, const Par_Vector& x) const;
	void Restriction(int level, const Par_Vector& b, const Par_Vector& x, const Par_Vector& g) const;
	void Prolongation(int level, const Par_Vector& e, const Par_Vector& x) const;
	void V_Cycle(int level, const Par_Vector& b, const Par_Vector& x) const;
	void W_Cycle(int level, const Par_Vector& b, const Par_Vector& x) const;
	void F_Cycle(int level, const Par_Vector& b, const Par_Vector& x) const;

public:
	int    MaxLevels;
	int    CoarseSize;
	double StrengthThreshold;
	int    AggNumLevels;
	int    CoarsenType;
	int    InterpMinElements;
	int    InterpMaxElements;
	double TruncationFactor;
	double SparsificationThreshold;
	int    CycleType;
	int    PreSweeps;
	int    PostSweeps;
	int    CoarseSweeps;
	int    PrintStats;

	ParBSR_ILUAMG_Preconditioner();
	ParBSR_ILUAMG_Preconditioner(
	int    max_levels,
	int    coarse_size,
	double strength_threshold,
	int    agg_num_levels,
	int    coarsen_type,
	int    interp_min_elements,
	int    interp_max_elements,	
	double truncation_factor,
	double sparsification_threshold,
	int    cycle_type,
	int    pre_sweeps,
	int    post_sweeps,
	int    coarse_sweeps,
	int    print_stats);
	~ParBSR_ILUAMG_Preconditioner();
	
	void Free();
	void Setup(const ParBSR_Matrix& A, int REUSE = 0);
	int InSize() const;
	int OutSize() const;
	void Apply(const Par_Vector& b, const Par_Vector& x) const;
};

}

#endif